import numpy as N
import scipy as S
import pylab as P

def shift(data, sh):
    data = N.hstack([data[sh:],data[:sh]])
    return data

def gauss_fourier(len_N,freq,exp=1):
    t = N.arange(len_N)
    sigma = 1 / float(N.abs(freq)**exp)
    w = 1/(sigma * N.sqrt(2*N.pi)) * N.exp(-1 * t**2 / float(2 * sigma**2))
    gauss_w = N.fft.fftshift(N.fft.fft(w))
    return gauss_w

def GetGFTIndices(length):
    signal_freq = \
    N.fft.fftshift(N.fft.fftfreq(length,1/float(length))).astype(int)
    result=[(0,0,1)]
    exp_2 = int(N.ceil(N.log2(length)))
    f_start = N.hstack([0,2**N.arange(exp_2-1)])
    for f_now in f_start:
        for sign in [1,-1]:
            if N.abs(f_now) < 2:
                f_width = 1
                f_centre = f_now + 1
            else:
                f_width = f_now
                f_centre = f_now + f_width/2 + 1
            f_end = f_now + f_width
            f_now *= sign
            f_centre *= sign
            f_end *= sign
            if sign == 1:
                if f_end > signal_freq.max(): f_end = signal_freq.max()
                result.append((f_now,f_centre,f_end-f_now))
            if sign == -1:
                if f_end < signal_freq.min(): f_end = signal_freq.min()
                result.append((f_end,f_centre,f_now-f_end))
    return N.array(result)

def GetGFTCoords(length):
    indices = GetGFTIndices(length)
    signal_freq = \
    N.fft.fftshift(N.fft.fftfreq(length,1/float(length))).astype(int)
    f = [0]
    t = [int(S.where(signal_freq==0)[0])]
    for i,centre in enumerate(indices[1:,1]):
        width = int(indices[i+1,2])
        f.extend(centre * N.ones([width],dtype=int))
        t.extend(N.round(N.linspace(int(length/float(width*2)), length-int(length/float(width*2)), width, \
                            endpoint=True),decimals=0).astype(int))
    coords = N.empty([length,],dtype=[('f',int,(1)), ('t',int,(1))])
    f = N.array(f)
    t = N.array(t)
    coords['f'] = f
    coords['t'] = t
    #order = coords.argsort(order=['f','t'])
    return coords #, order



def GFT_ST(data):
    len_N = len(data)
    signal_fft = N.fft.fftshift(N.fft.fft(data))
    signal_freq = N.fft.fftshift(N.fft.fftfreq(len_N,1/float(len_N)).astype(int))
    A_tau = [N.fft.ifftshift(signal_fft)[0]]
    gft_indices = GetGFTIndices(len_N)
    for (f_now, f_centre, f_width) in gft_indices[1:]:
        f_end = f_now + f_width
        index_1 = S.where(signal_freq == f_now)[0]
        index_2 = S.where(signal_freq == f_end)[0]
        shift_amount = S.where(signal_freq == 0)[0] - S.where(signal_freq == f_centre)[0]
        gauss_w = gauss_fourier(len_N,f_centre)
        gauss_w = shift(gauss_w,shift_amount)
        A_v = signal_fft[index_1:index_2] * gauss_w[index_1:index_2]
        A_tau.extend(N.fft.ifft(N.fft.ifftshift(A_v)))
    A_tau = N.array(A_tau)
    return A_tau

def Mesh_GFT(gft_result, kind='nearest'):
    len_N = len(gft_result)
    gft_indices = GetGFTIndices(len_N)
    signal_freq = N.fft.fftshift(N.fft.fftfreq(len_N,1/float(len_N)).astype(int))
    gft_type = gft_result.dtype
    if kind == 'nearest':
        mesh = N.empty([len_N,len_N],dtype=gft_type)
        k=0
        for i,width in enumerate(gft_indices[:,2]):
            index_1 = S.where(signal_freq == gft_indices[i,0])[0]
            index_2 = S.where(signal_freq == gft_indices[i,0] + \
                              gft_indices[i,2])[0]
            mesh[index_1:index_2,:] = \
            gft_result[N.round(N.linspace(k,k+width-1,len_N, endpoint=True), \
            decimals=0).astype(int)] * N.ones(width)[:,N.newaxis]
            k += width
    elif kind == 'spline':
        coords = GetGFTCoords(len_N)
        from scipy.interpolate import SmoothBivariateSpline
        f_interp = \
                SmoothBivariateSpline(y=coords['f'],x=coords['t'],z=gft_result.real)
        mesh_real = f_interp(y=signal_freq,x=N.arange(1024))
        f_interp = \
                SmoothBivariateSpline(y=coords['f'],x=coords['t'],z=gft_result.imag)
        mesh_imag = f_interp(y=signal_freq,x=N.arange(1024))
        mesh = mesh_real + 1j*mesh_imag
    elif kind in ['linear', 'cubic', 'quintic']:
        kx = ky = {'linear' : 1,
                   'cubic' : 3,
                   'quintic' : 5}[kind]
        s = len_N
        coords = GetGFTCoords(len_N)
        from scipy.interpolate import bisplrep, bisplev
        f_interp = \
                bisplrep(y=coords['f'],x=coords['t'],z=gft_result.real,kx=kx,ky=ky,s=s)
        1/0
        mesh_real = f_interp(y=signal_freq,x=N.arange(1024))
        f_interp = \
                bisplrep(y=coords['f'],x=coords['t'],z=gft_result.imag,kx=kx,ky=ky,s=s)
        mesh_imag = f_interp(y=signal_freq,x=N.arange(1024))
        mesh = mesh_real + 1j*mesh_imag
    return mesh

t = N.linspace(-128,128,900,endpoint=False) # Generate time variable
test_signal = S.where(N.abs(t) <=20, 1., 0.).astype(float)
test_signal = test_signal - N.mean(test_signal)
gft_result = GFT_ST(test_signal)

gft_mesh = N.abs(Mesh_GFT(gft_result,kind='nearest'))
gft_mmean = N.max(gft_mesh,axis=1)
gft_norm = S.where(gft_mmean[:,N.newaxis]!= 0,gft_mesh/gft_mmean[:,N.newaxis],N.zeros([len(t)]))
signal_freq = N.fft.fftshift(N.fft.fftfreq(len(t),1/float(5000)))

